import sys

from keras import regularizers

sys.path.append("/home/zxh/otu_classifier/")
import pickle

import keras
from keras.layers import Dense, Dropout, Activation, Flatten, Conv2D, MaxPooling2D
from keras.layers.normalization import BatchNormalization
from keras.models import Sequential
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import EarlyStopping
from keras.optimizers import SGD, Adam
from keras.utils import to_categorical
from keras.losses import sparse_categorical_crossentropy

from src.config import params, config
from src.datas.segmentationLabelsAndSample import labelAndSample

hps = params.get_default_params()
nb_epoch = 1
root_path = '../../model'


class Model:
    def __init__(self, X_train, X_test, y_train, y_test):
        self._X_train = X_train
        self._X_test = X_test
        self._y_train = y_train
        self._y_test = y_test
        self.model = None

    def build_model(self):
        self.model = Sequential()

        self.model.add(Conv2D(32, (1, 1), strides=1, padding='same',
                              input_shape=(self._X_train.shape[1], self._X_train.shape[2], self._X_train.shape[3])))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))

        self.model.add(Conv2D(32, (5, 5), padding='same'))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))

        self.model.add(Conv2D(32, (3, 3), padding='same'))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))

        self.model.add(Conv2D(64, (5, 5), padding='same'))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))

        self.model.add(Flatten())
        self.model.add(Dense(2048))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))

        self.model.add(Dropout(0.8))
        self.model.add(Dense(1024))
        self.model.add(BatchNormalization())
        self.model.add(Activation('relu'))

        self.model.add(Dropout(0.5))
        self.model.add(Dense(hps.num_classes))

        self.model.add(Activation('softmax'))
        self.model.summary()

    def train_model(self):
        adm = Adam(lr=hps.learning_rate, epsilon=1e-08)  # 设置动态变化的学习率
        self.model.compile(loss='categorical_crossentropy',
                           optimizer=adm,
                           metrics=['accuracy'])
        self.model.fit(self._X_train, self._y_train, batch_size=hps.batch_size, epochs=params.num_train_steps,
                       validation_data=(self._X_test, self._y_test))
        score = model.evaluate(self._X_test, self._y_test)
        print('model trained')

    def save_model(self):
        model_json = self.model.to_json()
        with open(root_path + "/model_json.json", "w") as json_file:
            json_file.write(model_json)
        self.model.save_weights(root_path + '/model_weight.h5')
        self.model.save(root_path + '/model.h5')
        print('model saved')


if __name__ == '__main__':
    X_train, y_train = pickle.load(open(config.train_pkl, 'rb'))
    X_test, y_test = pickle.load(open(config.test_pkl, 'rb'))
    y_train_lables, y_train_sample = labelAndSample(y_train)
    y_test_lables, y_test_sample = labelAndSample(y_test)
    y_train = to_categorical(y_train_lables, 3)
    y_test = to_categorical(y_test_lables, 3)
    model = Model(X_train, X_test, y_train, y_test)
    model.build_model()
    print('model built')
    model.train_model()
    print('model trained')
    model.save_model()
    print('model saved')
